package com.atomgit.ai.proxy.processor;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.atomgit.ai.proxy.entity.ModelInfo;
import com.atomgit.ai.proxy.entity.UserApiKey;
import com.atomgit.ai.proxy.param.PromptParam;
import com.atomgit.ai.proxy.param.inner.qianfan.QianFanAnswer;
import com.atomgit.ai.proxy.param.inner.qianfan.QianFanParam;
import com.atomgit.ai.proxy.param.inner.qianfan.QianFanParamMessage;
import com.atomgit.ai.proxy.param.inner.qianfan.QianFanParamParameters;
import com.atomgit.ai.proxy.processor.service.ParamAndResponseProcessor;
import com.atomgit.ai.proxy.vo.ModelAnswerOutPutVO;
import com.atomgit.ai.proxy.vo.ModelAnswerVO;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.*;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.RequestCallback;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.util.UriComponentsBuilder;

import javax.servlet.http.HttpServletResponse;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.*;

@Slf4j
public class ParamAndResponseProcessorForQianFan implements ParamAndResponseProcessor {


    @Override
    public String getApiKey(ModelInfo modelInfo,UserApiKey userApiKey) {
        return null;
    }

    @Override
    public String getUrl(ModelInfo modelInfo,UserApiKey userApiKey) {
        return String.format("%s?access_token=%s", modelInfo.getRequestUrl(),getAccessToken(
                modelInfo.getAccessTokenUrl(),
                HttpMethod.POST,
                new HashMap<String, String>(){{
                    put("client_id", userApiKey.getApiKey());
                    put("client_secret",userApiKey.getSecretKey());
                    put("grant_type","client_credentials");
                }}
        ));
    }


    public String getAccessToken(String url, HttpMethod httpMethod, Map<String, String> queryParameters) {
        try{
            RestTemplate restTemplate = new RestTemplate();
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            HttpEntity<String> requestEntity = new HttpEntity<>(headers);
            UriComponentsBuilder builder = UriComponentsBuilder
                    .fromHttpUrl(url);
            queryParameters.forEach(builder::queryParam);
            ResponseEntity<HashMap> response = restTemplate.exchange(builder.build().toUriString(),httpMethod,requestEntity, HashMap.class);
            return response.getBody().get("access_token").toString();
        }catch (Exception e){
            e.printStackTrace();
        }
        return null;
    }

    @Override
    public String generateRequestBody(PromptParam param) {
        List<QianFanParamMessage> qianFanParamMessages = new ArrayList<>();

        if (param.getHistory() != null)
            param.getHistory().forEach(answerHistory -> {
                qianFanParamMessages.add(new QianFanParamMessage(){{
                    setRole("user");
                    setContent(answerHistory.getUser());
                }});
                qianFanParamMessages.add(new QianFanParamMessage(){{
                    setRole("assistant");
                    setContent(answerHistory.getBot());
                }});
            });
        qianFanParamMessages.add(new QianFanParamMessage(){{
            setRole("user");
            setContent(param.getPrompt());
        }});
        QianFanParamParameters qianFanParamParameters = new QianFanParamParameters();
        // TODO: 2024/7/16 完善参数拷贝
        QianFanParam qianFanParam = new QianFanParam(){{
            setMessages(qianFanParamMessages);
            setQianFanParamParameters(qianFanParamParameters);
        }};
        return JSON.toJSONString(qianFanParam);
    }

    @Override
    public ModelAnswerVO generateModelAnswerVO(String responseBody) {
        responseBody = responseBody.replace("NULL", "");
        if (responseBody != null) {
            try {
                // 处理错误
                JSONObject jsonObject = JSON.parseObject(responseBody);
                if (jsonObject.containsKey("error_code") && jsonObject.getString("error_code")!=null&&!jsonObject.getString("error_code").equals("0")&&!jsonObject.getString("error_code").trim().equals("")){
                    return new ModelAnswerVO(){{
                        setOutput(new ModelAnswerOutPutVO(){{
                            setFinish_reason("stop");
                            setText(jsonObject.getString("error_msg"));
                        }});
                    }};
                }

                QianFanAnswer qianFanAnswer = JSON.parseObject(responseBody, QianFanAnswer.class);
                return new ModelAnswerVO(){{
                    setOutput(new ModelAnswerOutPutVO(){{
                        setFinish_reason(qianFanAnswer.getFinish_reason());
                        setText(qianFanAnswer.getResult());
                    }});
                }};
            }catch (Exception e){
                e.printStackTrace();
                return new ModelAnswerVO(){{
                    setOutput(new ModelAnswerOutPutVO(){{
                        setFinish_reason("stop");
                        setText("目前出现了一点错误");
                    }});
                }};
            }

        }
        return null;
    }


    public String sendRequestInStream(String url, HttpMethod httpMethod, String requestBody, String apiKey, HttpServletResponse response, ParamAndResponseProcessor paramAndResponseProcessor) {
        log.info("sendRequestInStream url:{},httpMethod:{},requestBody:{}",url,httpMethod,requestBody);
        try {
            RestTemplate restTemplate = new RestTemplate();
            // 设置请求头
            MultiValueMap<String, String> headersParam = new HttpHeaders();
            HttpHeaders headers = new HttpHeaders(headersParam);
            headers.setAccept(Collections.singletonList(MediaType.APPLICATION_JSON));
            headers.setContentType(MediaType.APPLICATION_JSON);
            // 添加其他需要的头部信息
            headers.add("Authorization", apiKey);
            headers.add("Accept", "text/event-stream");
            RequestCallback requestCallback = request -> {
                // 设置请求体
                request.getHeaders().addAll(headers);
                request.getBody().write(requestBody.getBytes(StandardCharsets.UTF_8));
            };
            response.setCharacterEncoding("UTF-8");
            restTemplate.execute(URI.create(url), httpMethod, requestCallback, clientHttpResponse -> {
                try (BufferedReader reader = new BufferedReader(new InputStreamReader(clientHttpResponse.getBody(), StandardCharsets.UTF_8))) {
                    // 检车http返回码是否为401(apikey/token失效)
                    if (clientHttpResponse.getStatusCode().value() == 401) {
                        log.info("sendRequest response:{}",response);
                        return "Invalid API-key provided.";
                    }
                    String line;
                    while ((line = reader.readLine()) != null) {
                        ModelAnswerVO modelAnswerVO = paramAndResponseProcessor.generateModelAnswerVO(line);
                        response.getWriter().println(JSON.toJSONString(modelAnswerVO));
                        response.getWriter().flush();
                    }
                    return null;
                }
            });
        } catch (Exception e) {
            log.info(e.getMessage());
            e.printStackTrace();
        }
        return null;
    }

}
